from langchain.agents import ConversationalAgent, AgentExecutor
from langchain.memory import ConversationBufferMemory  # 对话记忆组件
from langchain.tools import tool
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI


# 复用之前的工具（天气、计算器、知识库）
@tool(
    description="用于查询指定城市的天气。参数:city(必填,城市名称,示例:北京,苏州,上海),返回值:实时温度、湿度、天气状况、风向风力。用户问城市天气的时候必须调用这个工具,不可编造数据"
)
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": "温度：20℃，湿度：58%，天气：多云，风向：西风，风力：1级",
        "上海": "温度：24℃，湿度：50%，天气：晴，风向：南风，风力：2级",
        "广州": "温度：30℃，湿度：70%，天气：晴，风向：东风，风力：3级",
    }
    # 使用模糊匹配，处理可能的空格或其他字符问题
    city = city.strip()  # 去除首尾空格
    for key, value in weather_data.items():
        if key in city or city in key:
            return value
    return "没有该城市的天气信息"


@tool(
    description="用于计算加减乘除运算.参数:expression(必填,数学表达式,示例: 2+4,10-2,3*6,10/2,2+3*7),返回值:计算结果.仅处理数字和+、-、*、/"
)
def CalculatorTool(expression: str) -> str:
    res = eval(expression)
    return f"计算结果:{expression}={res}"


@tool(
    description="用于查询公司知识库中的答案。参数:question(必填,问题,示例:中国有多少个民族),返回值:答案。用户问相关问题时必须调用这个工具,不可编造数据"
)
def RAGTool(question: str) -> str:
    rag_data = {
        "中国有多少个民族": "中国有56个民族",
    }
    for key, value in rag_data.items():
        if key in question:
            return value
    return f"在知识库中没有找到[{question}]答案"


tools = [WeatherTool, CalculatorTool, RAGTool]

# 1. 初始化LLM（和之前一样）
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key="sk-6434b554122c4b049ceb805d703f695b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 2. 核心：创建对话记忆（存储历史对话）
memory = ConversationBufferMemory(
    memory_key="chat_history",  # 记忆在prompt中的变量名
    return_messages=True  # 返回Message对象，方便Agent理解
)

# 3. 构建适用于对话的提示模板（必须包含chat_history）
prompt_template = """
你是一个能进行多轮对话的智能助手，拥有以下工具:
{tools}

使用规则:
1. 先理解问题，结合对话历史{chat_history}判断是否需要工具；
2. 调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3. 拿到工具结果后，判断是否需要继续调用工具，否则整理结果回答用户。

用户的问题:{input}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "input", "chat_history"],  # 新增chat_history变量
)

# 4. 构建CONVERSATIONAL_REACT Agent（核心区别：用ConversationalAgent）
agent = ConversationalAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True,  # 打印思考过程
)

# 5. 构建执行器（传入记忆组件）
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,  # 关键：将记忆绑定到执行器
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
)

if __name__ == '__main__':
    # 多轮对话示例
    print("===== 第一轮对话 =====")
    print("对话历史:")
    print(memory.buffer)
    res1 = agent_executor.invoke({"input": "北京今天的天气怎么样？"})
    print("回答：", res1["output"], "\n")

    print("===== 第二轮对话 =====")
    print("对话历史:")
    print(memory.buffer)
    res2 = agent_executor.invoke({"input": "那适合穿什么衣服？"})  # 依赖上轮的“北京天气”
    print("回答：", res2["output"], "\n")

    print("===== 第三轮对话 =====")
    print("对话历史:")
    print(memory.buffer)
    res3 = agent_executor.invoke({"input": "如果买一件外套（299元）和一条裤子（150元），一共多少钱？"})
    print("回答：", res3["output"], "\n")
